#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
🚀 GPU加速定日镜场优化求解器
利用CuPy进行大规模并行计算，实现高性能优化
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime
import os
import time
import warnings
warnings.filterwarnings('ignore')

# 标准库
from scipy.optimize import differential_evolution
from scipy.spatial.distance import cdist

# GPU加速库
try:
    import cupy as cp
    import cupyx.scipy.optimize as cp_opt
    import cupyx.scipy.spatial.distance as cp_distance
    # 测试GPU是否真的可用
    test_array = cp.array([1, 2, 3])
    GPU_AVAILABLE = True
    print(f"✅ GPU加速可用: CuPy {cp.__version__}")
    try:
        device_name = cp.cuda.Device().name
        print(f"✅ GPU设备: {device_name}")
    except:
        print("✅ GPU设备: 未知")
except ImportError as e:
    print(f"❌ CuPy不可用: {e}")
    import scipy.optimize as sp_opt
    import scipy.spatial.distance as sp_distance
    GPU_AVAILABLE = False
except Exception as e:
    print(f"❌ GPU初始化失败: {e}")
    import scipy.optimize as sp_opt
    import scipy.spatial.distance as sp_distance
    GPU_AVAILABLE = False

class GPUHeliostatOptimizer:
    """GPU加速的定日镜场优化器"""
    
    def __init__(self, use_gpu=True):
        self.use_gpu = use_gpu and GPU_AVAILABLE
        self.xp = cp if self.use_gpu else np
        
        # 物理常数
        self.DNI = 800.0  # W/m²
        self.TOWER_HEIGHT = 100.0  # m
        self.MIRROR_SIZE = 6.0  # m²
        
        # 优化参数
        self.n_mirrors = 632  # 镜子数量
        self.field_radius = 350.0  # 镜场半径
        
        print(f"🔧 优化器初始化完成 ({'GPU' if self.use_gpu else 'CPU'}模式)")
    
    def cosine_efficiency_gpu(self, positions, tower_pos, sun_vector):
        """GPU加速余弦效率计算"""
        if self.use_gpu:
            positions = cp.asarray(positions)
            tower_pos = cp.asarray(tower_pos)
            sun_vector = cp.asarray(sun_vector)
        
        # 计算镜子到塔的向量
        mirror_to_tower = tower_pos - positions
        distances = self.xp.linalg.norm(mirror_to_tower, axis=1)
        mirror_to_tower_unit = mirror_to_tower / distances[:, None]
        
        # 计算镜面法向量（理想跟踪）
        normal_vectors = sun_vector + mirror_to_tower_unit
        normal_vectors = normal_vectors / self.xp.linalg.norm(normal_vectors, axis=1)[:, None]
        
        # 余弦效率
        cosine_eff = self.xp.abs(self.xp.dot(normal_vectors, sun_vector))
        
        return cosine_eff if self.use_gpu else cosine_eff
    
    def atmospheric_attenuation_gpu(self, distances):
        """GPU加速大气衰减计算"""
        if self.use_gpu:
            distances = cp.asarray(distances)
        
        # 简化的大气衰减模型
        visibility = 23.0  # km
        beta = 3.91 / visibility  # 衰减系数
        attenuation = self.xp.exp(-beta * distances / 1000.0)
        
        return attenuation
    
    def shadowing_blocking_gpu(self, positions, sun_vector):
        """GPU加速阴影遮挡计算"""
        n_mirrors = len(positions)
        
        if self.use_gpu:
            positions = cp.asarray(positions)
            sun_vector = cp.asarray(sun_vector)
            
            # 使用GPU并行计算距离矩阵
            distances = cp_distance.cdist(positions, positions)
            
            # 创建遮挡掩码
            blocking_mask = cp.zeros(n_mirrors, dtype=cp.float32)
            
            # 简化的遮挡计算（并行化）
            mirror_radius = cp.sqrt(self.MIRROR_SIZE / cp.pi)
            for i in range(n_mirrors):
                # 计算其他镜子对当前镜子的遮挡
                other_positions = positions[cp.arange(n_mirrors) != i]
                if len(other_positions) > 0:
                    # 简化遮挡判断
                    dist_to_others = distances[i, cp.arange(n_mirrors) != i]
                    potential_blockers = cp.sum(dist_to_others < 2 * mirror_radius)
                    blocking_mask[i] = cp.maximum(0, 1 - potential_blockers * 0.05)
            
            return blocking_mask
        else:
            # CPU版本
            blocking_efficiency = np.ones(n_mirrors)
            mirror_radius = np.sqrt(self.MIRROR_SIZE / np.pi)
            
            for i in range(n_mirrors):
                # 简化计算
                other_positions = np.delete(positions, i, axis=0)
                if len(other_positions) > 0:
                    distances = np.linalg.norm(other_positions - positions[i], axis=1)
                    close_mirrors = np.sum(distances < 2 * mirror_radius)
                    blocking_efficiency[i] = max(0, 1 - close_mirrors * 0.05)
            
            return blocking_efficiency
    
    def intercept_efficiency_gpu(self, positions, tower_pos):
        """GPU加速截取效率计算"""
        if self.use_gpu:
            positions = cp.asarray(positions)
            tower_pos = cp.asarray(tower_pos)
        
        distances = self.xp.linalg.norm(positions - tower_pos, axis=1)
        
        # 截取效率模型（简化）
        receiver_radius = 5.0  # m
        mirror_error = 0.002  # rad
        
        beam_radius = distances * mirror_error + np.sqrt(self.MIRROR_SIZE / np.pi)
        intercept_eff = self.xp.minimum(1.0, (receiver_radius / beam_radius) ** 2)
        
        return intercept_eff
    
    def total_efficiency_gpu(self, positions, sun_azimuth=180, sun_elevation=45):
        """GPU加速总效率计算"""
        tower_pos = self.xp.array([0, 0, self.TOWER_HEIGHT])
        
        # 太阳矢量
        sun_az_rad = self.xp.radians(sun_azimuth)
        sun_el_rad = self.xp.radians(sun_elevation)
        sun_vector = self.xp.array([
            self.xp.cos(sun_el_rad) * self.xp.sin(sun_az_rad),
            self.xp.cos(sun_el_rad) * self.xp.cos(sun_az_rad),
            self.xp.sin(sun_el_rad)
        ])
        
        # 各种效率计算
        cosine_eff = self.cosine_efficiency_gpu(positions, tower_pos, sun_vector)
        
        distances = self.xp.linalg.norm(positions - tower_pos, axis=1)
        atm_eff = self.atmospheric_attenuation_gpu(distances)
        
        shadow_eff = self.shadowing_blocking_gpu(positions, sun_vector)
        intercept_eff = self.intercept_efficiency_gpu(positions, tower_pos)
        
        # 总效率
        total_eff = cosine_eff * atm_eff * shadow_eff * intercept_eff
        
        return total_eff
    
    def objective_function(self, positions_flat):
        """优化目标函数"""
        positions = positions_flat.reshape(-1, 3)
        
        # 确保位置在有效范围内
        positions[:, 2] = 0  # 地面高度
        
        # 计算总效率
        efficiency = self.total_efficiency_gpu(positions)
        
        if self.use_gpu:
            efficiency = cp.asnumpy(efficiency)
        
        # 总功率（负值用于最大化）
        total_power = -self.xp.sum(efficiency) * self.DNI * self.MIRROR_SIZE
        
        return float(total_power)
    
    def generate_initial_layout(self):
        """生成初始镜场布局"""
        positions = []
        
        # 使用改进的径向布局
        n_rings = 8
        mirrors_per_ring = self.n_mirrors // n_rings
        
        for ring in range(1, n_rings + 1):
            ring_radius = (ring / n_rings) * self.field_radius
            n_mirrors_in_ring = min(mirrors_per_ring, 
                                  self.n_mirrors - len(positions))
            
            for i in range(n_mirrors_in_ring):
                angle = 2 * np.pi * i / n_mirrors_in_ring
                # 添加随机扰动
                angle += np.random.normal(0, 0.1)
                radius = ring_radius + np.random.normal(0, 10)
                
                x = radius * np.cos(angle)
                y = radius * np.sin(angle)
                
                # 避免距离塔太近
                if np.sqrt(x**2 + y**2) > 50:
                    positions.append([x, y, 0])
        
        # 补充不足的镜子
        while len(positions) < self.n_mirrors:
            angle = np.random.uniform(0, 2*np.pi)
            radius = np.random.uniform(80, self.field_radius)
            x = radius * np.cos(angle)
            y = radius * np.sin(angle)
            positions.append([x, y, 0])
        
        return np.array(positions[:self.n_mirrors])
    
    def optimize_layout(self, max_iterations=50):
        """优化镜场布局"""
        print("🚀 开始GPU加速优化...")
        start_time = time.time()
        
        # 生成初始布局
        initial_positions = self.generate_initial_layout()
        print(f"📍 初始布局生成完成: {len(initial_positions)} 个定日镜")
        
        # 计算初始性能
        initial_efficiency = self.total_efficiency_gpu(initial_positions)
        if self.use_gpu:
            initial_efficiency = cp.asnumpy(initial_efficiency)
        
        initial_power = np.sum(initial_efficiency) * self.DNI * self.MIRROR_SIZE
        print(f"🔋 初始总功率: {initial_power/1e6:.2f} MW")
        
        # 设置优化边界
        bounds = []
        for i in range(self.n_mirrors):
            bounds.extend([(-self.field_radius, self.field_radius),  # x
                          (-self.field_radius, self.field_radius),  # y
                          (0, 0)])  # z固定为0
        
        # 使用差分进化算法优化
        print("🎯 开始差分进化优化...")
        
        result = differential_evolution(
            self.objective_function,
            bounds,
            maxiter=max_iterations,
            popsize=10,
            seed=42,
            disp=True,
            workers=1  # GPU计算不需要多进程
        )
        
        optimization_time = time.time() - start_time
        print(f"⏱️ 优化完成，耗时: {optimization_time:.2f} 秒")
        
        # 获取优化结果
        optimized_positions = result.x.reshape(-1, 3)
        optimized_efficiency = self.total_efficiency_gpu(optimized_positions)
        if self.use_gpu:
            optimized_efficiency = cp.asnumpy(optimized_efficiency)
        
        optimized_power = np.sum(optimized_efficiency) * self.DNI * self.MIRROR_SIZE
        
        improvement = (optimized_power - initial_power) / initial_power * 100
        
        print(f"✅ 优化结果:")
        print(f"   🔋 优化后总功率: {optimized_power/1e6:.2f} MW")
        print(f"   📈 性能提升: {improvement:.2f}%")
        print(f"   🎯 平均效率: {np.mean(optimized_efficiency):.4f}")
        
        return {
            'initial_positions': initial_positions,
            'optimized_positions': optimized_positions,
            'initial_power': initial_power,
            'optimized_power': optimized_power,
            'improvement': improvement,
            'optimization_time': optimization_time,
            'initial_efficiency': initial_efficiency,
            'optimized_efficiency': optimized_efficiency
        }
    
    def analyze_performance(self, results):
        """性能分析"""
        print("\n📊 详细性能分析:")
        
        initial_eff = results['initial_efficiency']
        optimized_eff = results['optimized_efficiency']
        
        print(f"效率统计:")
        print(f"  初始: 平均={np.mean(initial_eff):.4f}, 标准差={np.std(initial_eff):.4f}")
        print(f"  优化后: 平均={np.mean(optimized_eff):.4f}, 标准差={np.std(optimized_eff):.4f}")
        
        # 效率分布分析
        high_eff_initial = np.sum(initial_eff > 0.5)
        high_eff_optimized = np.sum(optimized_eff > 0.5)
        
        print(f"高效率镜子数量 (>0.5):")
        print(f"  初始: {high_eff_initial}/{len(initial_eff)} ({high_eff_initial/len(initial_eff)*100:.1f}%)")
        print(f"  优化后: {high_eff_optimized}/{len(optimized_eff)} ({high_eff_optimized/len(optimized_eff)*100:.1f}%)")
    
    def visualize_results(self, results):
        """可视化结果"""
        plt.style.use('seaborn-v0_8')
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        
        initial_pos = results['initial_positions']
        optimized_pos = results['optimized_positions']
        initial_eff = results['initial_efficiency']
        optimized_eff = results['optimized_efficiency']
        
        # 1. 初始布局
        scatter1 = axes[0,0].scatter(initial_pos[:, 0], initial_pos[:, 1], 
                                   c=initial_eff, cmap='viridis', s=20)
        axes[0,0].set_title('初始布局')
        axes[0,0].set_xlabel('X (m)')
        axes[0,0].set_ylabel('Y (m)')
        axes[0,0].axis('equal')
        plt.colorbar(scatter1, ax=axes[0,0], label='效率')
        
        # 2. 优化布局
        scatter2 = axes[0,1].scatter(optimized_pos[:, 0], optimized_pos[:, 1], 
                                   c=optimized_eff, cmap='viridis', s=20)
        axes[0,1].set_title('优化后布局')
        axes[0,1].set_xlabel('X (m)')
        axes[0,1].set_ylabel('Y (m)')
        axes[0,1].axis('equal')
        plt.colorbar(scatter2, ax=axes[0,1], label='效率')
        
        # 3. 效率对比
        x = np.arange(len(initial_eff))
        axes[1,0].plot(x, initial_eff, 'o-', label='初始', alpha=0.6, markersize=2)
        axes[1,0].plot(x, optimized_eff, 'o-', label='优化后', alpha=0.6, markersize=2)
        axes[1,0].set_title('效率对比')
        axes[1,0].set_xlabel('镜子编号')
        axes[1,0].set_ylabel('效率')
        axes[1,0].legend()
        axes[1,0].grid(True)
        
        # 4. 效率分布直方图
        axes[1,1].hist(initial_eff, bins=30, alpha=0.5, label='初始', density=True)
        axes[1,1].hist(optimized_eff, bins=30, alpha=0.5, label='优化后', density=True)
        axes[1,1].set_title('效率分布')
        axes[1,1].set_xlabel('效率')
        axes[1,1].set_ylabel('密度')
        axes[1,1].legend()
        axes[1,1].grid(True)
        
        plt.tight_layout()
        
        # 保存图片
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"GPU优化结果_{timestamp}.png"
        plt.savefig(filename, dpi=300, bbox_inches='tight')
        print(f"📊 可视化结果已保存: {filename}")
        
        return fig
    
    def save_results(self, results):
        """保存结果到Excel"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"GPU优化结果_{timestamp}.xlsx"
        
        with pd.ExcelWriter(filename, engine='openpyxl') as writer:
            # 优化后位置
            pos_df = pd.DataFrame(results['optimized_positions'], 
                                columns=['X (m)', 'Y (m)', 'Z (m)'])
            pos_df['效率'] = results['optimized_efficiency']
            pos_df['功率 (W)'] = results['optimized_efficiency'] * self.DNI * self.MIRROR_SIZE
            pos_df.to_excel(writer, sheet_name='优化布局', index_label='镜子编号')
            
            # 性能总结
            summary_df = pd.DataFrame({
                '参数': ['初始总功率 (MW)', '优化后总功率 (MW)', '性能提升 (%)', 
                        '优化时间 (秒)', '镜子数量', '平均效率'],
                '数值': [results['initial_power']/1e6, results['optimized_power']/1e6,
                        results['improvement'], results['optimization_time'],
                        len(results['optimized_positions']), 
                        np.mean(results['optimized_efficiency'])]
            })
            summary_df.to_excel(writer, sheet_name='性能总结', index=False)
        
        print(f"💾 结果已保存: {filename}")
        return filename

def main():
    """主函数"""
    print("🚀 GPU加速定日镜场优化器启动")
    print("=" * 50)
    
    # 创建优化器
    optimizer = GPUHeliostatOptimizer(use_gpu=True)
    
    # 性能基准测试
    if optimizer.use_gpu:
        print("🔬 GPU性能基准测试...")
        positions = cp.random.rand(1000, 3) * 300
        
        # 测试GPU计算速度
        start_time = time.time()
        for _ in range(10):
            eff = optimizer.total_efficiency_gpu(positions)
        gpu_time = time.time() - start_time
        print(f"⚡ GPU计算时间: {gpu_time:.4f} 秒")
    
    # 执行优化
    results = optimizer.optimize_layout(max_iterations=30)
    
    # 分析性能
    optimizer.analyze_performance(results)
    
    # 可视化结果
    optimizer.visualize_results(results)
    
    # 保存结果
    optimizer.save_results(results)
    
    print("\n✅ 优化完成!")

if __name__ == "__main__":
    main() 